core/iter/adapters/
take.rs

1use crate::cmp;
2use crate::iter::adapters::SourceIter;
3use crate::iter::{FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, TrustedRandomAccess};
4use crate::num::NonZero;
5use crate::ops::{ControlFlow, Try};
6
7/// An iterator that only iterates over the first `n` iterations of `iter`.
8///
9/// This `struct` is created by the [`take`] method on [`Iterator`]. See its
10/// documentation for more.
11///
12/// [`take`]: Iterator::take
13/// [`Iterator`]: trait.Iterator.html
14#[derive(Clone, Debug)]
15#[must_use = "iterators are lazy and do nothing unless consumed"]
16#[stable(feature = "rust1", since = "1.0.0")]
17pub struct Take<I> {
18    iter: I,
19    n: usize,
20}
21
22impl<I> Take<I> {
23    pub(in crate::iter) const fn new(iter: I, n: usize) -> Take<I> {
24        Take { iter, n }
25    }
26}
27
28#[stable(feature = "rust1", since = "1.0.0")]
29impl<I> Iterator for Take<I>
30where
31    I: Iterator,
32{
33    type Item = <I as Iterator>::Item;
34
35    #[inline]
36    fn next(&mut self) -> Option<<I as Iterator>::Item> {
37        if self.n != 0 {
38            self.n -= 1;
39            self.iter.next()
40        } else {
41            None
42        }
43    }
44
45    #[inline]
46    fn nth(&mut self, n: usize) -> Option<I::Item> {
47        if self.n > n {
48            self.n -= n + 1;
49            self.iter.nth(n)
50        } else {
51            if self.n > 0 {
52                self.iter.nth(self.n - 1);
53                self.n = 0;
54            }
55            None
56        }
57    }
58
59    #[inline]
60    fn count(mut self) -> usize {
61        if self.n == 0 {
62            return 0;
63        }
64        // Advancing consumes the same elements `next` would have yielded,
65        // while benefiting from the inner iterator's `advance_by` fast path.
66        match self.iter.advance_by(self.n) {
67            Ok(()) => self.n,
68            Err(remaining) => self.n - remaining.get(),
69        }
70    }
71
72    #[inline]
73    fn size_hint(&self) -> (usize, Option<usize>) {
74        if self.n == 0 {
75            return (0, Some(0));
76        }
77
78        let (lower, upper) = self.iter.size_hint();
79
80        let lower = cmp::min(lower, self.n);
81
82        let upper = match upper {
83            Some(x) if x < self.n => Some(x),
84            _ => Some(self.n),
85        };
86
87        (lower, upper)
88    }
89
90    #[inline]
91    fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
92    where
93        Fold: FnMut(Acc, Self::Item) -> R,
94        R: Try<Output = Acc>,
95    {
96        fn check<'a, T, Acc, R: Try<Output = Acc>>(
97            n: &'a mut usize,
98            mut fold: impl FnMut(Acc, T) -> R + 'a,
99        ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a {
100            move |acc, x| {
101                *n -= 1;
102                let r = fold(acc, x);
103                if *n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) }
104            }
105        }
106
107        if self.n == 0 {
108            try { init }
109        } else {
110            let n = &mut self.n;
111            self.iter.try_fold(init, check(n, fold)).into_try()
112        }
113    }
114
115    #[inline]
116    fn fold<B, F>(self, init: B, f: F) -> B
117    where
118        Self: Sized,
119        F: FnMut(B, Self::Item) -> B,
120    {
121        Self::spec_fold(self, init, f)
122    }
123
124    #[inline]
125    fn for_each<F: FnMut(Self::Item)>(self, f: F) {
126        Self::spec_for_each(self, f)
127    }
128
129    #[inline]
130    #[rustc_inherit_overflow_checks]
131    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
132        let min = self.n.min(n);
133        let rem = match self.iter.advance_by(min) {
134            Ok(()) => 0,
135            Err(rem) => rem.get(),
136        };
137        let advanced = min - rem;
138        self.n -= advanced;
139        NonZero::new(n - advanced).map_or(Ok(()), Err)
140    }
141}
142
143#[unstable(issue = "none", feature = "inplace_iteration")]
144unsafe impl<I> SourceIter for Take<I>
145where
146    I: SourceIter,
147{
148    type Source = I::Source;
149
150    #[inline]
151    unsafe fn as_inner(&mut self) -> &mut I::Source {
152        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
153        unsafe { SourceIter::as_inner(&mut self.iter) }
154    }
155}
156
157#[unstable(issue = "none", feature = "inplace_iteration")]
158unsafe impl<I: InPlaceIterable> InPlaceIterable for Take<I> {
159    const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY;
160    const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY;
161}
162
163#[stable(feature = "double_ended_take_iterator", since = "1.38.0")]
164impl<I> DoubleEndedIterator for Take<I>
165where
166    I: DoubleEndedIterator + ExactSizeIterator,
167{
168    #[inline]
169    fn next_back(&mut self) -> Option<Self::Item> {
170        if self.n == 0 {
171            None
172        } else {
173            let n = self.n;
174            self.n -= 1;
175            self.iter.nth_back(self.iter.len().saturating_sub(n))
176        }
177    }
178
179    #[inline]
180    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
181        let len = self.iter.len();
182        if self.n > n {
183            let m = len.saturating_sub(self.n) + n;
184            self.n -= n + 1;
185            self.iter.nth_back(m)
186        } else {
187            if len > 0 {
188                self.iter.nth_back(len - 1);
189            }
190            None
191        }
192    }
193
194    #[inline]
195    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
196    where
197        Self: Sized,
198        Fold: FnMut(Acc, Self::Item) -> R,
199        R: Try<Output = Acc>,
200    {
201        if self.n == 0 {
202            try { init }
203        } else {
204            let len = self.iter.len();
205            if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
206                try { init }
207            } else {
208                self.iter.try_rfold(init, fold)
209            }
210        }
211    }
212
213    #[inline]
214    fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
215    where
216        Self: Sized,
217        Fold: FnMut(Acc, Self::Item) -> Acc,
218    {
219        if self.n == 0 {
220            init
221        } else {
222            let len = self.iter.len();
223            if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
224                init
225            } else {
226                self.iter.rfold(init, fold)
227            }
228        }
229    }
230
231    #[inline]
232    #[rustc_inherit_overflow_checks]
233    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
234        // The amount by which the inner iterator needs to be shortened for it to be
235        // at most as long as the take() amount.
236        let trim_inner = self.iter.len().saturating_sub(self.n);
237        // The amount we need to advance inner to fulfill the caller's request.
238        // take(), advance_by() and len() all can be at most usize, so we don't have to worry
239        // about having to advance more than usize::MAX here.
240        let advance_by = trim_inner.saturating_add(n);
241
242        let remainder = match self.iter.advance_back_by(advance_by) {
243            Ok(()) => 0,
244            Err(rem) => rem.get(),
245        };
246        let advanced_by_inner = advance_by - remainder;
247        let advanced_by = advanced_by_inner - trim_inner;
248        self.n -= advanced_by;
249        NonZero::new(n - advanced_by).map_or(Ok(()), Err)
250    }
251}
252
253#[stable(feature = "rust1", since = "1.0.0")]
254impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {}
255
256#[stable(feature = "fused", since = "1.26.0")]
257impl<I> FusedIterator for Take<I> where I: FusedIterator {}
258
259#[unstable(issue = "none", feature = "trusted_fused")]
260unsafe impl<I: TrustedFused> TrustedFused for Take<I> {}
261
262#[unstable(feature = "trusted_len", issue = "37572")]
263unsafe impl<I: TrustedLen> TrustedLen for Take<I> {}
264
265trait SpecTake: Iterator {
266    fn spec_fold<B, F>(self, init: B, f: F) -> B
267    where
268        Self: Sized,
269        F: FnMut(B, Self::Item) -> B;
270
271    fn spec_for_each<F: FnMut(Self::Item)>(self, f: F);
272}
273
274impl<I: Iterator> SpecTake for Take<I> {
275    #[inline]
276    default fn spec_fold<B, F>(mut self, init: B, f: F) -> B
277    where
278        Self: Sized,
279        F: FnMut(B, Self::Item) -> B,
280    {
281        use crate::ops::NeverShortCircuit;
282        self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
283    }
284
285    #[inline]
286    default fn spec_for_each<F: FnMut(Self::Item)>(mut self, f: F) {
287        // The default implementation would use a unit accumulator, so we can
288        // avoid a stateful closure by folding over the remaining number
289        // of items we wish to return instead.
290        fn check<'a, Item>(
291            mut action: impl FnMut(Item) + 'a,
292        ) -> impl FnMut(usize, Item) -> Option<usize> + 'a {
293            move |more, x| {
294                action(x);
295                more.checked_sub(1)
296            }
297        }
298
299        let remaining = self.n;
300        if remaining > 0 {
301            self.iter.try_fold(remaining - 1, check(f));
302        }
303    }
304}
305
306impl<I: Iterator + TrustedRandomAccess> SpecTake for Take<I> {
307    #[inline]
308    fn spec_fold<B, F>(mut self, init: B, mut f: F) -> B
309    where
310        Self: Sized,
311        F: FnMut(B, Self::Item) -> B,
312    {
313        let mut acc = init;
314        let end = self.n.min(self.iter.size());
315        for i in 0..end {
316            // SAFETY: i < end <= self.iter.size() and we discard the iterator at the end
317            let val = unsafe { self.iter.__iterator_get_unchecked(i) };
318            acc = f(acc, val);
319        }
320        acc
321    }
322
323    #[inline]
324    fn spec_for_each<F: FnMut(Self::Item)>(mut self, mut f: F) {
325        let end = self.n.min(self.iter.size());
326        for i in 0..end {
327            // SAFETY: i < end <= self.iter.size() and we discard the iterator at the end
328            let val = unsafe { self.iter.__iterator_get_unchecked(i) };
329            f(val);
330        }
331    }
332}
333
334#[stable(feature = "exact_size_take_repeat", since = "1.82.0")]
335impl<T: Clone> DoubleEndedIterator for Take<crate::iter::Repeat<T>> {
336    #[inline]
337    fn next_back(&mut self) -> Option<Self::Item> {
338        self.next()
339    }
340
341    #[inline]
342    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
343        self.nth(n)
344    }
345
346    #[inline]
347    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
348    where
349        Self: Sized,
350        Fold: FnMut(Acc, Self::Item) -> R,
351        R: Try<Output = Acc>,
352    {
353        self.try_fold(init, fold)
354    }
355
356    #[inline]
357    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
358    where
359        Self: Sized,
360        Fold: FnMut(Acc, Self::Item) -> Acc,
361    {
362        self.fold(init, fold)
363    }
364
365    #[inline]
366    #[rustc_inherit_overflow_checks]
367    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
368        self.advance_by(n)
369    }
370}
371
372// Note: It may be tempting to impl DoubleEndedIterator for Take<RepeatWith>.
373// One must fight that temptation since such implementation wouldn’t be correct
374// because we have no way to return value of nth invocation of repeater followed
375// by n-1st without remembering all results.
376
377#[stable(feature = "exact_size_take_repeat", since = "1.82.0")]
378impl<T: Clone> ExactSizeIterator for Take<crate::iter::Repeat<T>> {
379    fn len(&self) -> usize {
380        self.n
381    }
382}
383
384#[stable(feature = "exact_size_take_repeat", since = "1.82.0")]
385impl<F: FnMut() -> A, A> ExactSizeIterator for Take<crate::iter::RepeatWith<F>> {
386    fn len(&self) -> usize {
387        self.n
388    }
389}
core/iter/adapters/take.rs

core/iter/adapters/
take.rs
